import numpy as np
import mindspore as ms
from skimage.metrics import peak_signal_noise_ratio, structural_similarity

# compute the PSNR and SSIM between the orginal image and the result image(get by learning)
def get_PSNR_SSIM(orginal_clean_image, result_image, data_range):
    # convert to numpy array
    if not isinstance(orginal_clean_image, np.ndarray):
        if isinstance(orginal_clean_image, ms.Tensor):
            orginal_clean_image = orginal_clean_image.asnumpy()
    if not isinstance(result_image, np.ndarray):
        if isinstance(result_image, ms.Tensor):
            result_image = result_image.asnumpy()
    image_num = orginal_clean_image[0]
    PSNR = []
    SSIM = []
    for i in range(image_num):
        PSNR.append(peak_signal_noise_ratio(orginal_clean_image[i, :, :, :],
                                            result_image[i, :, :, :],
                                            data_range=data_range))
        orginal_clean_image_ = np.expand_dims(np.squeeze(orginal_clean_image), axis = -1)
        result_image_ = np.expand_dims(np.squeeze(result_image), axis = -1)
        SSIM.append(structural_similarity(np.array(orginal_clean_image_, dtype=np.float16),
                                          np.array(result_image_, dtype=np.float16),
                                          data_range=data_range, multichannel=True))
    return np.mean(PSNR), np.mean(SSIM)